//!  fftw3  傅里叶变换c2c-----------------------------------
//输入complex 数据进行 fft 和 ifft 计算并返回complex数据类型
#ifndef FOURIER_H_
#define FOURIER_H_
#include <fftw3.h>
#include <iostream>
#include <cmath>
#include <complex>
#define complex std::complex<double>

namespace Fourier{
    void fft(int N, complex* in, complex* out); //对complex 数据进行fft
    void ifft(int N, complex* in, complex* out);//对complex 数据进行ifft，获得complex数组

    void fft_r2c(int N, double* in, complex* out);//对double 数据进行fft
    void ifft_c2r(int N, complex* in, double* out);//对complex 数据进行ifft ,获得double 数据

    void fftfreq(int N, double* array, double timestep);
}

void Fourier::fft(int N, complex* in, complex* out){ //对complex 数据进行fft
        fftw_plan p_forward;
        fftw_complex* m_in, * m_out;
        m_in = (fftw_complex*)fftw_malloc(sizeof(fftw_complex) * N);
        m_out = (fftw_complex*)fftw_malloc(sizeof(fftw_complex) * N);

        //cout << "输入数据" <<endl;
        for(int i = 0; i < N ;i++){
            m_in[i][0] = real(in[i]);
            m_in[i][1] = imag(in[i]);
            m_out[i][0]=0.;
            m_out[i][1]=0.;
        }

        //cout << "执行变换"<< endl;
        p_forward = fftw_plan_dft_1d(N, m_in, m_out, FFTW_FORWARD, FFTW_ESTIMATE);       
        fftw_execute(p_forward);    /* repeat as needed */

        for (int i = 0; i < N; i++){
            out[i].real(m_out[i][0]);
            out[i].imag(m_out[i][1]);
            //cout << out_fft[i].real() << " + i*" << out[i].imag()<< endl;
        }
        //cout<<"数据销毁"<<endl;
        fftw_destroy_plan(p_forward);
        fftw_free(m_in);
        fftw_free(m_out);
        //cout <<"out1地址:"<<out_fft<<endl;
        
}

void Fourier::ifft(int N, complex* in, complex* out){ //对complex 数据进行ifft，获得complex数组

        fftw_plan p_backward;
        fftw_complex* m_in, * m_out;
        m_in = (fftw_complex*)fftw_malloc(sizeof(fftw_complex) * N);
        m_out = (fftw_complex*)fftw_malloc(sizeof(fftw_complex) * N);

        //cout << "输入数据" <<endl;
        for(int i = 0; i < N ;i++){
            m_in[i][0] = real(in[i]);
            m_in[i][1] = imag(in[i]);
            m_out[i][0] = 0.;
            m_out[i][1] = 0.;
            //cout << m_in[i][0] << "+i*" << m_in[i][1] <<endl;
        }

        //cout << "执行逆变换" <<endl;
        p_backward = fftw_plan_dft_1d(N, m_in, m_out, FFTW_BACKWARD, FFTW_ESTIMATE);
        fftw_execute(p_backward);
        

        for (int i = 0; i < N; i++){
            out[i].real(m_out[i][0]/(double)N);
            out[i].imag(m_out[i][1]/(double)N);
            //cout << out_c2c[i].real() << " + i*" << out_c2c.imag() << endl;
        }

        //cout<<"数据销毁"<<endl;
        fftw_destroy_plan(p_backward);
        fftw_free(m_in);
        fftw_free(m_out);
        //cout <<"out2地址:"<<out_c2c<<endl;
}

void Fourier::fft_r2c(int N, double* in, complex* out) {//对double 数据进行fft
    
    fftw_plan p_forward;
    double* m_in = NULL;
    fftw_complex* m_out = NULL;
    m_in = (double*)fftw_malloc(sizeof(double) * N);
    m_out = (fftw_complex*)fftw_malloc(sizeof(fftw_complex) * N);

    //cout << "输入数据" <<endl;
    for (int i = 0; i < N; i++) {
        
        m_in[i] = in[i];//初始化 输入输出 数据网格
        m_out[i][0]=0.;
        m_out[i][1]=0.;
    }
    
    //cout << "执行变换"<< endl;
    p_forward = FFTW3_H::fftw_plan_dft_r2c_1d(N, m_in, m_out, FFTW_ESTIMATE);
    fftw_execute(p_forward);    /* repeat as needed */
    
    for (int i = 0; i < N; i++) {
        out[i].real(m_out[i][0]);
        out[i].imag(m_out[i][1]);
        //cout << out_fft[i].real() << " + i*" << out[i].imag()<< endl;
    }

    //cout<<"数据销毁"<<endl;
    fftw_destroy_plan(p_forward);
    fftw_free(m_in);
    fftw_free(m_out);
    //cout <<"out_fft地址:"<<out_fft<<endl;

}

void Fourier::ifft_c2r(int N, complex* in, double* out) {//对complex 数据进行ifft ,获得double 数据;in数据物理长度为m.
    
    fftw_plan p_backward;
    fftw_complex* m_in = NULL;
    double* m_out = NULL;
 
    m_in = (fftw_complex*)fftw_malloc(sizeof(fftw_complex) * N);
    m_out = (double*)fftw_malloc(sizeof(double) * N);

    int m = (N % 2 == 0)? int(N/2)+1 : int((N-1)/2)+1;
    //cout << "输入数据" <<endl;
    for (int i = 0; i < N; i++) {
        m_in[i][0] =(i < m )? real(in[i]) : 0;
        m_in[i][1] =(i < m )? imag(in[i]) : 0;
        m_out[i]=0;
        //cout << m_in[i][0] << "+i*" << m_in[i][1] <<endl;
    }

    //cout << "执行逆变换" <<endl;
    p_backward = FFTW3_H::fftw_plan_dft_c2r_1d(N, m_in, m_out, FFTW_ESTIMATE);
    fftw_execute(p_backward);
    
    for (int i = 0; i < N; i++) {
        out[i] = m_out[i] / (double)N;
    }

    //cout<<"数据销毁"<<endl;
    fftw_destroy_plan(p_backward);
    fftw_free(m_in);
    fftw_free(m_out);
    //cout <<"out2地址:"<<out_c2r<<endl;
}

void Fourier::fftfreq(int N, double* array, double timestep){ //获得频率轴

    double df = 1.0/double(N * timestep) ;
    for(int i = 0; i  <N; i++){
        if(N % 2 == 0){
            array[i] = (i <= N / 2 - 1)? (i * df) : ((i - N)*df);    
        }
        
        else{
            array[i] = (i <= (N - 1) / 2)? (i * df) : ((i - N)*df);
        }
    }    


}
#endif